package edu.washington.ling.syntric.proc;


import java.util.ArrayList;

import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.washington.ling.syntric.io.ExampleDoc;
import edu.washington.ling.syntric.io.ExampleDocReader;
import edu.washington.ling.syntric.io.ExampleDocWriter;
import edu.washington.ling.syntric.io.Example;
import edu.washington.ling.syntric.io.TreeReader;
import edu.washington.ling.syntric.io.TreeRepresentation;
import edu.washington.ling.syntric.io.TreeWriter;
import edu.washington.ling.syntric.struct.error.PrintingExceptionHandler;
import edu.washington.ling.syntric.struct.error.UnreadableRepresentationException;
import edu.washington.ling.syntric.struct.tree.Node;
import edu.washington.ling.syntric.struct.tree.Token;
import edu.washington.ling.syntric.struct.tree.Tree;
import edu.washington.ling.syntric.util.MergeUtils;
import edu.washington.ling.syntric.util.TreeUtils;


public class AttachTracesToDS {
	private static AttachTracesToDSOptions opts = new AttachTracesToDSOptions();
	
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
	
	private static void attachTrace(Token t, Tree tree) {
		Node adjNode = TreeUtils.findFirstBranching(t);
		ArrayList<String> headWords = new ArrayList<String>(adjNode.getHeadWords());
		if (headWords.size() == 1) {
			String headWord = headWords.get(0);
			if (tree.hasTokenWithName(headWord)) {
				t.setDSHead(tree.getTokenWithName(headWord));
				t.setDSRelation("ec");
			}
		}
		// Throw an error if this fails
	}
	private static void attachTraces(Tree tree) {
		for (Token t: tree.getTokenList()) {
			if (t.isPSEmptyCategory()) {
				attachTrace(t, tree);
			}
		}
	}
	
	private static void attachTracesForFile(){
		int numExamples = 0;
		int numWithBoth = 0;
		int numDSConstCreated = 0;
		try {
			ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, "xml");
			ExampleDoc exf = exfr.read();
			ExampleDocWriter exfw = ExampleDocWriter.newWriter(opts.outputFilePath, "xml");
			TreeReader dsr = TreeReader.newReader(opts.dsTreeFormat);
			TreeReader psr = TreeReader.newReader(opts.psTreeFormat);
			TreeWriter trw = TreeWriter.newWriter(opts.dsTreeFormat);
			Example exr = null;
			while ((exr = exf.getNextExample()) != null) {
				numExamples++;
				if (exr.hasTreeRep("DS") && exr.hasTreeRep("PS")) {
					numWithBoth++;
					TreeRepresentation dsrep = exr.getTreeRep("DS");
					try {
						Tree ds = dsr.readTree(dsrep);
						Tree ps = psr.readTree(exr.getTreeRep("PS"));
						Tree tree = MergeUtils.mergeDSandPSTrees(ds, ps, true, false);
						tree.setLayer("DS-Const");
						TreeUtils.labelHeadWordsForTree(tree);
						attachTraces(tree);
						// nullify PS root to remove the PS component of 'tree'
						tree.setPSRoot(null);
						exr.addTreeRep(trw.writeTree(tree));
						numDSConstCreated++;
					}
					catch (UnreadableRepresentationException e) {
						String errStr = String.format("Could not process representation '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
				}
				
				exfw.writeExample(exr);
			}
			exfw.close();
		} 
		catch(Exception e) {
			pHandler.handle(e, "Attach PS traces to DS failed.", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
		System.out.println("Attach PS traces to DS was successful.");
		System.out.println(String.format("%d/%d DS const trees were created", numDSConstCreated, numWithBoth));
		System.out.println(String.format("Processed %d total examples", numExamples));
	}
	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			attachTracesForFile();
		}
		catch(CmdLineException e) {
			System.out.println(e.getMessage());
			System.out.print("Usage: AttachTracesToDS");
			parser.printSingleLineUsage(System.out);
			System.out.println();
			parser.printUsage(System.out);
			System.exit(-1);
		}
	}
}
